import fiftyone as fo
from fiftyone import ViewField as F

def image_directory_to_fiftyone(images_dir, dataset_name):
    """
    导入本地路径图片到fiftyone中
    """
    dataset = fo.Dataset.from_images_dir(
        dataset_dir=images_dir,
        dataset_type=fo.types.ImageDirectory,
        name=dataset_name,
        persistent=True,
    )
    print(f"导入了 {len(dataset)} 张图片")
    return dataset


def load_dataset_from_yolov5(dataset_name, dataset_dir=None, is_exist=False):
    """
    导入yolov5格式的标注数据到fiftyone中
    """
    if is_exist:
        dataset = fo.load_dataset(dataset_name)
        dataset.compute_metadata()
        print(f"加载了 {len(dataset)} 张图片")
        print(dataset.first())
        return dataset
    
    if dataset_dir is None:
        raise ValueError("dataset_dir must be provided")

    # The splits to load
    splits = ["train", "val"]

    # Load the dataset, using tags to mark the samples in each split
    dataset = fo.Dataset(dataset_name, overwrite=False)
    for split in splits:
        dataset.add_dir(
            dataset_dir=dataset_dir,
            dataset_type=fo.types.YOLOv5Dataset,
            split=split,
            tags=split,
        )

    # View summary info about the dataset
    # print(dataset)
    # Print the first few samples in the dataset
    # print(dataset.head())
    
    return dataset


if __name__ == "__main__":
    dataset_dir = r"D:\fiftyone\datasets\sihe-collection\3211-spraying-detect"
    dataset_name = "3211-spraying"
    dataset = load_dataset_from_yolov5(dataset_name, dataset_dir, is_exist=False)
    
    # Bboxes are in [top-left-x, top-left-y, width, height] format
    bbox_area = F("bounding_box")[2] * F("bounding_box")[3]

    # Only includes predictions whose bounding boxes have an area of at
    # least 50% of the image, and only include samples with at least
    # one prediction after filtering
    large_boxes_view = dataset.filter_labels("ground_truth", bbox_area >= 0.5)

    session = fo.launch_app(view=large_boxes_view)
    # session = fo.launch_app(dataset)  # （可选）启动FiftyOne App查看数据集
    session.wait()
